{
 "metadata": {
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.6"
  },
  "orig_nbformat": 2,
  "kernelspec": {
   "name": "python386jvsc74a57bd0304dd9ef76ad2d6bf7237ef8f6fde70bef676e0aff923029a3f8742854b86f8f",
   "display_name": "Python 3.8.6 64-bit ('normal': virtualenvwrapper)"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2,
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "code1 = '''int main()\n",
    "{\n",
    "\tint n, a, i;\n",
    "\tscanf(\"%d\",&n);\n",
    "\tfor (i = 1;i <= n;i++)\n",
    "\t{\n",
    "\t\tscanf(\"%d\",&a);\n",
    "\t\tprintf(\"%d\\n\",Factor(a, 2));\n",
    "\t}\n",
    "\treturn 0;\n",
    "}\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "code2 = '''\n",
    "int main()\n",
    "{\n",
    "\tint depart (int number, int y);\n",
    "\t\n",
    "\tint i, n, number;\n",
    "\tcin >> n;\n",
    "\tfor (i = 1; i <= n; i++)\n",
    "\t{\n",
    "\t\tcin >> number;\n",
    "\t\tcout << depart (number, 2) << endl;\n",
    "\t}\n",
    "\t\n",
    "\treturn 0;\n",
    "}\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "def check_folder(path):\n",
    "    if not Path(path).is_dir():\n",
    "        os.makedirs(path)\n",
    "def packcode(code1, code2):\n",
    "    def code2cpp(code, name):\n",
    "        cppfile = f'{name}.cpp'\n",
    "        check_folder(name)\n",
    "        with open(name + '/' + cppfile, 'w') as f:\n",
    "            f.write(code)\n",
    "    code2cpp(code1, 'code1')\n",
    "    code2cpp(code2, 'code2')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "packcode(code1,code2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil\n",
    "shutil.rmtree('a/b')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "0-cpg.dot\n"
     ]
    }
   ],
   "source": [
    "alldots = list(Path('dot/code1').glob('*.dot'))\n",
    "str(alldots[0].parent)\n",
    "print(alldots[0].name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "def get_main_dict(graphdir):\n",
    "    mainname = \"\"\n",
    "    alldots = list(Path(graphdir).glob('*.dot'))\n",
    "    r = '^digraph\\s(main)\\s{'\n",
    "    for dot in alldots:\n",
    "        path = str(dot.parent)\n",
    "        name = dot.name\n",
    "        #print(name)\n",
    "        class_id = path.split('\\\\')[-1]\n",
    "        t = str(name.split('-')[0])\n",
    "        with dot.open() as f:\n",
    "            cpg = f.readline()\n",
    "        #print(cpg)\n",
    "        if re.match(r, cpg):\n",
    "            mainname = name\n",
    "        return mainname"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_main_dot():\n",
    "    code1 = 'dot/code1'\n",
    "    code2 = 'dot/code2'\n",
    "    code1main = get_main_dict(code1)\n",
    "    code2main = get_main_dict(code2)\n",
    "    shutil.copyfile(code1 + '/' + code1main, './code1.dot')\n",
    "    shutil.copyfile(code2 + '/' + code2main, './code2.dot')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "get_main_dot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import subprocess\n",
    "import shutil\n",
    "import re\n",
    "import json\n",
    "import pygraphviz as pgv\n",
    "import networkx as nx\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "\n",
    "JOERNPATH = '/home/hls/joern-cli/'\n",
    "task = 'cpg14'\n",
    "\n",
    "\n",
    "def check_folder(path):\n",
    "    if not Path(path).is_dir():\n",
    "        os.makedirs(path)\n",
    "\n",
    "\n",
    "def del_file(path):\n",
    "    if Path(path).is_dir():\n",
    "        shutil.rmtree(path)\n",
    "\n",
    "\n",
    "def packcode(code1, code2):\n",
    "    def code2cpp(code, name):\n",
    "        cppfile = f'{name}.cpp'\n",
    "        check_folder(name)\n",
    "        with open(name + '/' + cppfile, 'w') as f:\n",
    "            f.write(code)\n",
    "    code2cpp(code1, 'code1')\n",
    "    code2cpp(code2, 'code2')\n",
    "\n",
    "\n",
    "def parse_source_code_to_dot():\n",
    "    # parse source code into cpg\n",
    "    code1 = 'code1'\n",
    "    code2 = 'code2'\n",
    "    check_folder('dot')\n",
    "\n",
    "    def parse(code, name):\n",
    "        print(f'parseing source code into {task}...')\n",
    "        shell_str = \"sh \" + JOERNPATH + \"joern-parse \" + code\n",
    "        subprocess.call(shell_str, shell=True)\n",
    "        print(f'exporting cpg from {task} root...')\n",
    "\n",
    "        outpath = f'dot/{name}'\n",
    "        del_file(outpath)\n",
    "        shell_export = \"sh \" + JOERNPATH + \"joern-export \" + \\\n",
    "            f\"--repr {task} --out \" + outpath\n",
    "        subprocess.call(shell_export, shell=True)\n",
    "    parse(code1, 'code1')\n",
    "    parse(code2, 'code2')\n",
    "# parse的dot文件分别存放在 dot/code1 dot/code2下\n",
    "# 获取main_dot\n",
    "def get_main_dict(graphdir):\n",
    "    mainname = \"\"\n",
    "    alldots = list(Path(graphdir).glob('*.dot'))\n",
    "    r = '^digraph\\s(main)\\s{'\n",
    "    for dot in alldots:\n",
    "        path = str(dot.parent)\n",
    "        name = dot.name\n",
    "        #print(name)\n",
    "        class_id = path.split('\\\\')[-1]\n",
    "        t = str(name.split('-')[0])\n",
    "        with dot.open() as f:\n",
    "            cpg = f.readline()\n",
    "        #print(cpg)\n",
    "        if re.match(r, cpg):\n",
    "            mainname = name\n",
    "        return mainname\n",
    "   \n",
    "def get_main_dot():\n",
    "    code1 = 'dot/code1'\n",
    "    code2 = 'dot/code2'\n",
    "    code1main = get_main_dict(code1)\n",
    "    code2main = get_main_dict(code2)\n",
    "    shutil.copyfile(code1 + '/' + code1main, './code1.dot')\n",
    "    shutil.copyfile(code2 + '/' + code2main, './code2.dot')\n",
    "\n",
    "def get_graph_massage(g, labeled=True):\n",
    "        \"\"\"\n",
    "        输入一张图，返回 edges 和 labels\n",
    "        \"\"\"\n",
    "        num_nodes = g.number_of_nodes()\n",
    "        num_edges = g.number_of_edges()\n",
    "        node_labels = [\"\" for i in range(num_nodes)]\n",
    "        edge_labels = []\n",
    "\n",
    "        # 给节点进行重新编号\n",
    "        id = 0\n",
    "        new_id = {}\n",
    "        # 建立node 的label和 node的id之间的映射\n",
    "        for node in g.nodes():\n",
    "            new_id[node] = id\n",
    "            id += 1\n",
    "        \n",
    "        if labeled == True:\n",
    "            hsh = list(g.nodes())\n",
    "            for node in g.nodes().data(): #形如 ('7', {'type': 'C', 'label': '7'})\n",
    "                # labels[7] = 8\n",
    "                ori_id = node[0]\n",
    "                now_id = new_id[ori_id]\n",
    "                if node[1] :\n",
    "                    #node_labels[now_id] = node[1][\"label\"]\n",
    "                    node_labels[now_id] = node[1]['label'].split(',')[0][1:]\n",
    "                    # for condition of \"<operator>.lessThan\"\n",
    "                    if '<operator>' in node_labels[now_id]:\n",
    "                        node_labels[now_id] = node_labels[now_id].split('.')[1]\n",
    "                else:\n",
    "                    node_labels[now_id] = \"@\"\n",
    "        edges = []\n",
    "        # edges 的每个元素形如： [7, 4]\n",
    "        for edge in g.edges().data(): #形如 ('7', '4', {'valence': 2, 'id': '6'})\n",
    "            edges.append([new_id[edge[0]], new_id[edge[1]]])\n",
    "            edge_labels.append(edge[2][\"label\"])\n",
    "        # 返回一个图的边矩阵, 以及顶点向量；\n",
    "        # 这是json的组成形式\n",
    "        nodes = list(new_id.values())\n",
    "        return nodes, edges, node_labels, edge_labels\n",
    "def getgraph(dotfile):\n",
    "    Gtmp1 = pgv.AGraph(dotfile)\n",
    "    g = nx.DiGraph(Gtmp1)\n",
    "    return g\n",
    "def getpair():\n",
    "    g1 = getgraph('code1.dot')\n",
    "    nodes1, edges1, nodes_label1, edges_label1 = get_graph_massage(g1, labeled=True)\n",
    "    g2 = getgraph('code2.dot')\n",
    "    nodes2, edges2, nodes_label2, edges_label2 = get_graph_massage(g2, labeled=True)\n",
    "    ret = {}\n",
    "    ret[\"graph1\"] = edges1\n",
    "    ret[\"graph2\"] = edges2\n",
    "    ret[\"nodes1\"] = nodes1\n",
    "    ret[\"nodes2\"] = nodes2\n",
    "    # 每个节点的文本内容\n",
    "    ret[\"nodes_label1\"] = nodes_label1\n",
    "    ret[\"nodes_label2\"] = nodes_label2\n",
    "    # 每条边的属性\n",
    "    ret[\"edges_label1\"] = edges_label1\n",
    "    ret[\"edges_label2\"] = edges_label2\n",
    "    with open('pair.json', 'w') as jf:\n",
    "        json.dump(ret, jf)\n",
    "\n",
    "def preprocess(code1, code2):\n",
    "    packcode(code1, code2)\n",
    "    parse_source_code_to_dot()\n",
    "    get_main_dot()\n",
    "    getpair()\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# if __name__ == '__main__':\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "preprocess()"
   ]
  }
 ]
}